#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AsyncAnalysisPass - 异步函数分析阶段

职责：
1. 标记async函数（IRFunctionNode.is_async = True）
2. 收集函数中的await表达式位置
3. 验证await只在async函数中使用
4. 分析控制流，检测不支持的模式（如异步循环）
5. 设置extension参数
"""

from typing import Any, Dict, Optional, List
from dataclasses import dataclass

from mcpy.ir.module import IRModule
from .base import BasePass, PassStage, PassContext
from ..ir.base import IRNode
from ..ir.declarations import IRFunctionNode, IRMethodNode
from ..ir.expressions.await_expr import IRAwait
from ..ir.statements import IRForRange, IRForEach, IRWhile


@dataclass
class AwaitInfo:
    """await表达式信息"""

    node: IRAwait
    line: int
    in_loop: bool = False  # 是否在循环中
    in_try: bool = False  # 是否在try块中


class AsyncAnalysisError(Exception):
    """异步分析错误"""

    pass


class AsyncAnalysisPass(BasePass):
    """异步函数分析Pass

    分析async函数和await表达式，设置必要的元数据。
    """

    name = "async_analysis"
    stage = PassStage.TYPE_ANALYSIS

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        super().__init__()
        self.config = config or {}
        self.errors: List[str] = []

    def run(self, ir_module: IRModule, context: PassContext) -> IRNode:
        """运行异步分析

        Args:
            ir_module: IR模块节点
            context: Pass上下文

        Returns:
            处理后的IR模块节点
        """
        self.errors.clear()

        # 遍历所有节点进行分析
        self._analyze_async_functions(ir_module)

        # 如果有错误，抛出异常
        if self.errors:
            error_msg = "\n".join(self.errors)
            raise AsyncAnalysisError(f"异步分析失败:\n{error_msg}")

        return ir_module

    def _analyze_async_functions(self, node: IRNode) -> None:
        """分析async函数

        Args:
            node: 要分析的节点
        """
        # 如果是函数节点且是async函数
        if isinstance(node, (IRFunctionNode, IRMethodNode)) and node.is_async:
            self._analyze_single_async_function(node)

        # 递归处理子节点
        for child in node.get_child_nodes():
            self._analyze_async_functions(child)

    def _analyze_single_async_function(self, func_node: IRFunctionNode) -> None:
        """分析单个async函数

        Args:
            func_node: async函数节点
        """
        # 收集await表达式
        await_exprs: List[AwaitInfo] = []
        self._collect_await_expressions(func_node, await_exprs, in_loop=False)

        # 为 await 表达式提取类型信息（用于 stackless 代码生成）
        self._extract_await_types(func_node, await_exprs)

        # 设置extension参数
        has_await = len(await_exprs) > 0

        # 分析闭包变量（用于 stackless 代码生成）
        if has_await:
            self._analyze_closure_variables(func_node)
        func_node.set_extension_param("has_await_expressions", has_await)
        func_node.set_extension_param("await_chain_info", await_exprs)
        func_node.set_extension_param(
            "needs_io_context", True
        )  # async函数需要io_context

        # 循环中的 await 现在已支持（通过递归控制流处理）

    def _analyze_closure_variables(self, func_node: IRFunctionNode) -> None:
        """分析函数中引用的闭包变量（同级嵌套函数）

        Args:
            func_node: 函数节点
        """
        from mcpy.ir.expressions.basic import IRName

        captured = set()

        def visit_node(node):
            """递归访问节点，查找外部变量引用"""
            if isinstance(node, IRName):
                # 查找符号
                symbol = node.lookup_symbol(node.id)
                if symbol and symbol.ir_node:
                    # 检查是否是嵌套函数
                    if isinstance(symbol.ir_node, IRFunctionNode):
                        # 查找被引用函数的外层函数
                        ref_outer_func = self._find_outer_function(symbol.ir_node)
                        # 查找当前函数的外层函数
                        cur_outer_func = self._find_outer_function(func_node)

                        if ref_outer_func is None:
                            # 顶层函数，全局可见，不需要捕获
                            pass
                        elif (
                            ref_outer_func == cur_outer_func
                            and symbol.ir_node != func_node
                        ):
                            # 同一个外层函数中的其他嵌套函数
                            # 检查是否是简单函数（无捕获），简单函数不需要捕获
                            ref_func_node = symbol.ir_node
                            is_simple = ref_func_node.get_extension_param(
                                "is_simple_non_capturing", False
                            )
                            # 如果标记未设置，检查 captured_vars 是否为空
                            if not is_simple:
                                captured_vars = getattr(ref_func_node, "captured_vars", None)
                                is_simple = not captured_vars or len(captured_vars) == 0

                            if not is_simple:
                                # 非简单函数（有捕获），需要捕获
                                captured.add(node.id)
                            # 简单函数不加入 captured，按普通作用域查找处理

            # 递归访问子节点
            for child in node.get_child_nodes():
                visit_node(child)

        # 遍历函数体
        visit_node(func_node.body)

        # 设置captured_vars
        if captured:
            func_node.captured_vars = list(captured)

    def _find_outer_function(self, func_node: IRFunctionNode):
        """查找函数的外层函数

        Args:
            func_node: 函数节点

        Returns:
            外层函数节点，如果没有返回 None
        """
        parent = func_node.parent
        while parent:
            if isinstance(parent, IRFunctionNode):
                return parent
            parent = parent.parent
        return None

    def _extract_await_types(
        self, func_node: IRFunctionNode, await_exprs: List[AwaitInfo]
    ) -> None:
        """提取 await 表达式的返回类型信息（Python 类型）

        为 IRAwait 节点设置 inferred_type（Python TypeInfo 对象）。
        这个方法工作在 Python 类型层级，不涉及 C++ 类型转换。

        Args:
            func_node: 函数节点
            await_exprs: await 表达式信息列表
        """
        for await_info in await_exprs:
            await_node = await_info.node
            await_expr = await_node.value  # await 的表达式

            # 提取 Python 类型
            python_type = self._extract_python_type_from_await_expr(await_expr)
            if python_type:
                # 设置 Python 类型（不是 cpp_type）
                await_node.set_extension_param("inferred_type", python_type)

    def _extract_python_type_from_await_expr(self, await_expr: IRNode):
        """从 await 表达式中提取 Python 类型

        Args:
            await_expr: 被 await 的表达式

        Returns:
            Python TypeInfo 对象，如果无法提取则返回 None
        """
        from mcpy.ir.expressions.call import IRCall
        from mcpy.ir.expressions.basic import IRName

        # 只处理函数调用的情况
        if not isinstance(await_expr, IRCall):
            return None

        if not isinstance(await_expr.func, IRName):
            return None

        # 查找被调用的函数
        called_func = self._lookup_called_function(await_expr.func)
        if not called_func:
            return None

        # 提取函数的返回类型（Python 类型）
        return self._extract_awaitable_inner_type(called_func)

    def _lookup_called_function(self, func_name_node: IRNode):
        """查找被调用的函数节点

        Args:
            func_name_node: 函数名节点（IRName）

        Returns:
            IRFunctionNode 或 None
        """
        from mcpy.ir.expressions.basic import IRName

        if not isinstance(func_name_node, IRName):
            return None

        called_symbol = func_name_node.lookup_symbol(func_name_node.id)
        if not called_symbol or not called_symbol.ir_node:
            return None

        from mcpy.ir.declarations.function import IRFunctionNode

        if isinstance(called_symbol.ir_node, IRFunctionNode):
            return called_symbol.ir_node

        return None

    def _extract_awaitable_inner_type(self, func_node: IRFunctionNode):
        """从异步函数的返回类型中提取内部类型

        例如：Awaitable[int] -> int

        Args:
            func_node: 函数节点

        Returns:
            内部类型（TypeInfo 对象）或 None
        """
        # 获取函数的 Python 返回类型
        inferred_return_type = func_node.get_extension_param("inferred_return_type")
        if not inferred_return_type:
            return None

        # 如果是 Awaitable[T]，提取 T
        if hasattr(inferred_return_type, "name") and inferred_return_type.name in (
            "Awaitable",
            "awaitable",
        ):
            if hasattr(inferred_return_type, "args") and inferred_return_type.args:
                return inferred_return_type.args[0]

        # 如果不是 Awaitable 类型，直接返回（可能是同步函数被误用）
        return inferred_return_type

    def _collect_await_expressions(
        self,
        node: IRNode,
        await_exprs: List[AwaitInfo],
        in_loop: bool = False,
        in_try: bool = False,
    ) -> None:
        """收集await表达式

        Args:
            node: 要检查的节点
            await_exprs: await表达式列表（输出参数）
            in_loop: 当前是否在循环中
            in_try: 当前是否在try块中
        """
        # 如果是await表达式
        if isinstance(node, IRAwait):
            await_info = AwaitInfo(
                node=node, line=node.source_line, in_loop=in_loop, in_try=in_try
            )
            await_exprs.append(await_info)

            # 继续检查await的值表达式
            if node.value:
                self._collect_await_expressions(
                    node.value, await_exprs, in_loop, in_try
                )
            return

        # 如果进入循环节点，设置in_loop标志
        is_loop_node = isinstance(node, (IRForRange, IRForEach, IRWhile))
        if is_loop_node:
            in_loop = True

        # 如果进入try节点，设置in_try标志
        from ..ir.statements import IRTryExcept

        if isinstance(node, IRTryExcept):
            in_try = True

        # 递归处理子节点
        for child in node.get_child_nodes():
            self._collect_await_expressions(child, await_exprs, in_loop, in_try)

    def _validate_await_usage(
        self, node: IRNode, current_function: Optional[IRFunctionNode] = None
    ) -> None:
        """验证await的使用是否合法

        Args:
            node: 要检查的节点
            current_function: 当前所在的函数
        """
        # 如果是函数节点，更新当前函数
        if isinstance(node, (IRFunctionNode, IRMethodNode)):
            current_function = node

        # 如果是await表达式
        if isinstance(node, IRAwait):
            # 检查是否在async函数中
            if not current_function or not current_function.is_async:
                self.errors.append(
                    f"第 {node.source_line} 行: await表达式只能在async函数中使用"
                )

        # 递归处理子节点
        for child in node.get_child_nodes():
            self._validate_await_usage(child, current_function)
